package com.github.uiautomator;

import android.annotation.SuppressLint;
import android.graphics.Bitmap;
import android.graphics.Matrix;
import android.graphics.Point;
import android.graphics.Rect;
import android.net.LocalServerSocket;
import android.net.LocalSocket;
import android.os.IBinder;
import android.os.Process;
import android.os.SystemClock;
import android.util.Log;

import com.github.uiautomator.compat.WindowManagerWrapper;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.lang.reflect.Method;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;

public class MinicapAgent extends Thread {
    private static final String SOCKET = "minicapagent";
    private static final String TAG = MinicapAgent.class.getSimpleName();
    private final WindowManagerWrapper windowManager = new WindowManagerWrapper();
    private int maxFPS;
    private int quality;
    private int size;
    private int width;
    private int height;
    private int rotation, hwrotate = 0;
    private Method screenshotMethod = null;
    private Method screenshotRectMethod = null;
    private boolean isRunning = false;

    public MinicapAgent(int size, int quality, int maxFPS) {
        this.size = size;
        this.quality = quality;
        this.maxFPS = maxFPS;
        rotation = windowManager.getRotation();
        Point p = windowManager.getDisplaySize();
        if(rotation % 2 == 0) {
            width = p.x;
            height = p.y;
        } else {
            width = p.y;
            height = p.x;
        }

        windowManager.watchRotation(r -> {
            rotation = r;

            Point p2 = windowManager.getDisplaySize();
            if(rotation % 2 == 0) {
                width = p2.x;
                height = p2.y;
            } else {
                width = p2.y;
                height = p2.x;
            }
            Log.v(TAG, "Rotation: " + r + ", width: " + width + ", height: " + height);

            isRunning = false;
        });

        try {
            Class surfaceControl = Class.forName("android.view.SurfaceControl");
            try {
                screenshotMethod = surfaceControl.getDeclaredMethod("screenshot", Integer.TYPE, Integer.TYPE);
                Log.i(TAG, "screenshot Int,Int");
            } catch(NoSuchMethodException e) {
                Log.e(TAG, "screenshot Int,Int:", e);
                try {
                    screenshotRectMethod = surfaceControl.getDeclaredMethod("screenshot", Rect.class, Integer.TYPE, Integer.TYPE, Integer.TYPE);
                    Log.i(TAG, "screenshot Rect,Int,Int,Int");
                } catch(NoSuchMethodException e2) {
                    Log.e(TAG, "screenshot Rect,Int,Int,Int:", e2);
                    Log.i(TAG, "captureDisplay");
                }
            }
        } catch(ClassNotFoundException e) {
            Log.e(TAG, "SurfaceControl.class:", e);
        }

        try {
            Class cls = Class.forName("android.os.SystemProperties");
            Method mth = cls.getMethod("getInt", String.class, Integer.TYPE);
            hwrotate = (int) mth.invoke(null, new Object[]{"persist.sf.hwrotation", 0});
        } catch (Exception e) {
            Log.e(TAG, "SystemProperties.getInt:", e);
            e.printStackTrace();
        }

        Log.i(TAG, "hwrotate = " + hwrotate);
    }

    private void manageClientConnection(LocalServerSocket serverSocket) {
        while (true) {
            try (LocalSocket socket = serverSocket.accept()) {
                Log.d(TAG, "client connected");

                OutputStream output = socket.getOutputStream();
                Log.v(TAG, "Display: " + width + "x" + height);
                Log.v(TAG, "Rotation: " + rotation);
                try {
                    writeBanner(output);
                    isRunning = true;
                    pipeImages(output);
                } catch (Exception e) {
                    Log.v(TAG, "socket closed, exception: " + e);
                } finally {
                    output.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    private void writeBanner(OutputStream stream) throws IOException {
        final byte BANNER_SIZE = 24;
        final byte version = 1;
        final byte quirks = 1; // QUIRK_DUMB
        int pid = Process.myPid();

        ByteBuffer b = ByteBuffer.allocate(BANNER_SIZE);
        b.order(ByteOrder.LITTLE_ENDIAN);
        b.put((byte) version); // version
        b.put(BANNER_SIZE);//banner size
        b.putInt(pid);//pid
        b.putInt(0);//real width
        b.putInt(0);//real height
        b.putInt(0);//desired width
        b.putInt(0);//desired height
        b.put((byte) rotation);//orientation
        b.put((byte) quirks);//quirks
        byte[] array = b.array();
        stream.write(array);
    }

    private void pipeImages(OutputStream stream) throws Exception {
        final long interval = 1000 / maxFPS;

        long count = 0;
        long lastStartTime = 0;
        long calulateStartTime = SystemClock.uptimeMillis();

        while (isRunning) {
            long waitMillis = lastStartTime + interval - SystemClock.uptimeMillis();
            if (waitMillis > 0) {
                Log.d(TAG, "wait time: " + waitMillis);
                SystemClock.sleep(waitMillis);
            }
            lastStartTime = SystemClock.uptimeMillis();
            Bitmap bmp = screenshot();
            try {
                long t = SystemClock.uptimeMillis();
                ByteArrayOutputStream buf = new ByteArrayOutputStream();
                bmp.compress(Bitmap.CompressFormat.JPEG, quality, buf);
                byte[] jpegData = buf.toByteArray();

                Log.d(TAG, "compress time: " + (SystemClock.uptimeMillis() - t));
                t = SystemClock.uptimeMillis();

                ByteBuffer b = ByteBuffer.allocate(4);
                b.order(ByteOrder.LITTLE_ENDIAN);
                b.putInt(jpegData.length);
                stream.write(b.array()); // write image size
                stream.write(jpegData);  // write image data

                Log.d(TAG, "send time: " + (SystemClock.uptimeMillis() - t));
                t = SystemClock.uptimeMillis();

                count++;
                if (count == maxFPS * 2) {
                    float fps = (float) count / (SystemClock.uptimeMillis() - calulateStartTime) * 1000;
                    Log.v(TAG, "FPS: " + fps);
                    calulateStartTime = SystemClock.uptimeMillis();
                    count = 0;
                }
            } finally {
                bmp.recycle();
            }
            Log.d(TAG, "run time: " + (SystemClock.uptimeMillis() - lastStartTime));
        }
    }

    private Bitmap screenshot() throws Exception {
        long t = SystemClock.uptimeMillis();
        Bitmap tmp = screenshotMethod != null ? takeScreenshot() : (screenshotRectMethod != null ? takeScreenshotRect() :  takeScreenshotDisplay());
        int width = tmp.getWidth();
        int height = tmp.getHeight();
        int w, h;

        Log.d(TAG, "screenshot time: " + (SystemClock.uptimeMillis() - t));

        t = SystemClock.uptimeMillis();

        Matrix matrix = new Matrix();

        if(width > size || height > size) {
            if (width > height) {
                if (width > size) {
                    w = size;
                    h = height * size / width;
                } else {
                    w = width;
                    h = height;
                }
            } else if (height > size) {
                h = size;
                w = width * size / height;
            } else {
                w = width;
                h = height;
            }

            matrix.postScale((float) w / width, (float) h / height);
            if(hwrotate == 90) {
                matrix.postRotate(90);
            }
        } else if(hwrotate == 90) {
            matrix.postRotate(90);
        } else {
            return tmp;
        }

        Bitmap bmp = Bitmap.createBitmap(tmp, 0, 0, width, height, matrix, false);

        tmp.recycle();

        Log.d(TAG, "resize time: " + (SystemClock.uptimeMillis() - t));

        return bmp;
    }

    private Bitmap takeScreenshot() throws Exception {
        try {
            Bitmap bmp = (Bitmap) screenshotMethod.invoke(null, new Object[]{width, height});
		    // Log.v(TAG, "[Default] width: " + bmp.getWidth() + ", height: " + bmp.getHeight() + ", rotation: " + rotation + ", w: " + width + ", h: " + height);
            if (rotation == 0) {
                return bmp;
            }
            long t = SystemClock.uptimeMillis();
            Matrix matrix = new Matrix();
            matrix.postRotate(-90 * this.rotation);
            Bitmap rotatedBmp = Bitmap.createBitmap(bmp, 0, 0, bmp.getWidth(), bmp.getHeight(), matrix, true);
            bmp.recycle();
            Log.d(TAG, "rotate time: " + (SystemClock.uptimeMillis() - t));
            return rotatedBmp;
        } catch (Exception e) {
            throw new Exception("Inject SurfaceControl fail", e);
        }
    }

    private Bitmap takeScreenshotRect() throws Exception {
        Bitmap bmp;
        if(hwrotate == 90) {
            bmp = (Bitmap) screenshotRectMethod.invoke(null, new Object[]{new Rect(0, 0, width, height), height, width, 0});
        } else {
            bmp = (Bitmap) screenshotRectMethod.invoke(null, new Object[]{new Rect(0, 0, width, height), width, height, rotation});
        }
	    return bmp;
    }

    @SuppressLint("BlockedPrivateApi")
    private Bitmap takeScreenshotDisplay() throws Exception {
    	int w, h;
    	Point size = windowManager.getDisplaySize();
        this.rotation = windowManager.getRotation();
        switch(this.rotation) {
        	case 0:
        	case 2:
        		w = size.x;
        		h = size.y;
        		break;
        	default:
        		w = size.y;
        		h = size.x;
        		break;
        }
    	
        Class surfaceControl = Class.forName("android.view.SurfaceControl");
        final long[] displayIds = (long[]) surfaceControl.getDeclaredMethod("getPhysicalDisplayIds").invoke(null, new Object[]{});
        final IBinder displayToken = (IBinder) surfaceControl.getDeclaredMethod("getPhysicalDisplayToken", Long.TYPE).invoke(null, new Object[]{displayIds[0]});
        Class Builder = Class.forName("android.view.SurfaceControl$DisplayCaptureArgs$Builder");
        Object builder = Builder.getConstructor(IBinder.class).newInstance(displayToken);
        Builder.getMethod("setSourceCrop", Rect.class).invoke(builder, new Object[]{new Rect(0, 0, w, h)});
        Builder.getMethod("setSize", Integer.TYPE, Integer.TYPE).invoke(builder, new Object[]{w, h});
        Object captureArgs = Builder.getMethod("build").invoke(builder, new Object[]{});
        Class DisplayCaptureArgs = Class.forName("android.view.SurfaceControl$DisplayCaptureArgs");
        Class ScreenshotHardwareBuffer = Class.forName("android.view.SurfaceControl$ScreenshotHardwareBuffer");
        Object screenshotBuffer = surfaceControl.getDeclaredMethod("captureDisplay", DisplayCaptureArgs).invoke(null, new Object[]{captureArgs});
        if(screenshotBuffer != null) {
            Bitmap bmp = (Bitmap) ScreenshotHardwareBuffer.getMethod("asBitmap").invoke(screenshotBuffer);
		    // Log.v(TAG, "[Display] width: " + bmp.getWidth() + ", height: " + bmp.getHeight() + ", rotation: " + rotation + ", w: " + w + ", h: " + h);
            return bmp;
        }
        return null;
    }

    @Override
    public void run() {
        try (LocalServerSocket serverSocket = new LocalServerSocket(SOCKET)) {
            Log.v(TAG, "Listening on localabstract:" + SOCKET);
            manageClientConnection(serverSocket);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
